#version 460 core

//模型数据
#define MESH_UNKNOWN 0
#define MESH_SPHERE 1
#define MESH_TRIANGLE 2
#define MESH_POINTLIGHT 3

//材质类型，0未知，1普通，2光源
#define MATERIAL_UNKNOWN 0
#define MATERIAL_NORMAL 1
#define MATERIAL_LIGHT 2
#define MATERIAL_REFLECT 3

const int Sample_Count = 1;     //采样数，即N*N次采样
const int MAX_DEPTH = 5;    //递归深度
uniform int traceDepth;         //当前光追深度
const float TRACE_COLOR_WEIGHT = 0.1f; //追踪颜色的权重
const float MAX_DISTANCE = 4;           //光最大距离

struct Ray{
    vec3 pos;
    vec3 dir;
};

struct RayHit
{
    vec3 pos;   //碰撞点
    vec3 ndir;   //碰撞点的法线
    vec3 reflectDir; //反射方向
    vec3 color; //碰撞点的颜色
    float material; //材质索引
    int isHit; //是否碰撞
};

struct MeshHead
{
    float type;     //类型
    float len;   //数据长度，是固定的
    float material; //材质索引
    float empty;    //留空
};

struct Sphere
{
    MeshHead head;
    vec3 center;
    float r;
    vec4 color;
};

struct Triangle
{
    MeshHead head;
    vec4 v[3];
    vec4 color;
};

struct PointLight
{
    MeshHead head;
    vec3 pos;
    float intensity;
    vec4 color;
};


struct RayTraceStackItem{
    vec3 lightColor;
    int index;
    RayHit hit;
};

uniform mat4 model;
uniform mat4 view;
uniform mat4 projection;
uniform sampler2D sceneData;
//uniform sampler2D noiseTex;

in vec3 pos;
out vec4 out_color;

//场景数据，以2D贴图
vec2 sceneDataSize;


vec2 IndexTo2D(float index){
    float y = int(index / sceneDataSize.x);
    float x = int(index) % int(sceneDataSize.x);
    return vec2(x,y);
}

//根据一维坐标获取场景数据
vec4 SceneDataAt(float index){
    vec2 uv = IndexTo2D(index);
    uv.x /= sceneDataSize.x;
    uv.y /= sceneDataSize.y;
    return texture(sceneData, uv);
}

//获取mesh头
MeshHead MeshHeadAt(float index){
    MeshHead ret;
    vec4 temp = SceneDataAt(index);
    ret.type = temp.r;
    ret.len = temp.g;
    ret.material = temp.b;
    return ret;
}


//获取一个球体
Sphere SphereAt(float index){
    Sphere ret;
    ret.head = MeshHeadAt(index);
    vec4 centerR = SceneDataAt(index+1);
    ret.center = centerR.xyz;
    ret.r = centerR.w;
    ret.color = SceneDataAt(index+2);
    return ret;
}

//获取一个三角形
Triangle TriangleAt(float index){
    Triangle ret;
    ret.head = MeshHeadAt(index);
    ret.v[0] = SceneDataAt(index+1);
    ret.v[1] = SceneDataAt(index+2);
    ret.v[2] = SceneDataAt(index+3);
    ret.color = SceneDataAt(index+4);
    return ret;
}

//获取一个点光源
PointLight PointLightAt(float index){
    PointLight ret;
    ret.head = MeshHeadAt(index);
    ret.pos = SceneDataAt(index+1).xyz;
    ret.intensity = SceneDataAt(index+1).w;
    ret.color = SceneDataAt(index+2);
    return ret;
}

//球体碰撞检测
RayHit SphereHit(Sphere s,Ray ray){
    RayHit ret;
    ret.isHit = 0;

    //求解射线与球体的交点        
    float shadowVal = dot(s.center - ray.pos,ray.dir);

    if(shadowVal < 0)
        return ret;

    vec3 shadowPos = ray.pos + shadowVal * ray.dir;

    if(length(shadowPos - s.center) >= s.r)
        return ret;

    
    //求最近的一个交点
    float h = length(shadowPos - s.center);

    float dis = sqrt(s.r * s.r - h * h);
    if(dis == 0) return ret;

    vec3 hitPos = shadowPos - dis * ray.dir;

    vec3 normal = normalize(hitPos - s.center);
    ret.pos = hitPos;
    ret.ndir = normal;
    ret.color = s.color.rgb;
    ret.material = s.head.material;
    ret.isHit = 1;

    return ret;
}


//三角形碰撞检测
RayHit TriangleHit(Triangle t,Ray ray){
    RayHit ret;
    ret.isHit = 0;

    vec3 p0 = t.v[0].xyz;
    vec3 p1 = t.v[1].xyz;
    vec3 p2 = t.v[2].xyz;

    vec3 v01 = p1 - p0;
    vec3 v02 = p2 - p0;

    vec3 n = cross(v01,v02);
    n = normalize(n);

    //射线与平面的交点
    float tval = dot(p0 - ray.pos,n) / dot(ray.dir,n);
    if(tval <= 0) return ret;

    vec3 hitPos = ray.pos + tval * ray.dir;

    //判断交点是否在三角形内部
    vec3 c0 = cross(p1 - p0,hitPos - p0);
    vec3 c1 = cross(p2 - p1,hitPos - p1);
    vec3 c2 = cross(p0 - p2,hitPos - p2);

    if(dot(c0,n) > 0 && dot(c1,n) > 0 && dot(c2,n) > 0)
    {
        ret.pos = hitPos;
        ret.ndir = n;
        ret.color = t.color.rgb;
        ret.material = t.head.material;
        ret.isHit = 1;
    }

    return ret;
}




//射线与世界的最近交点，只碰撞一次
RayHit worldHit(Ray ray){
    
    //最近的交点
    RayHit retHit;
    retHit.isHit = 0;

    if(ray.dir == vec3(0))
        return retHit;


    vec4 temp = SceneDataAt(0);
    float count = temp.r;
    float index = 1;


    for(int i = 0;i<count;i++)
    {
        MeshHead meshHead = MeshHeadAt(index);
        if(meshHead.len < 1) break;  //数据错误

        RayHit curHit;
        curHit.isHit = 0;

        if(meshHead.type == MESH_SPHERE)
            curHit = SphereHit(SphereAt(index),ray);
        else if(meshHead.type == MESH_TRIANGLE)
            curHit = TriangleHit(TriangleAt(index),ray);
       
        index += meshHead.len;

        //没有碰撞，跳过
        if(curHit.isHit == 0) continue;

        //距离太远跳过
        if(length(curHit.pos - ray.pos) > MAX_DISTANCE*2) continue;

        if(retHit.isHit == 0 || length(curHit.pos - ray.pos) < length(retHit.pos - ray.pos))
        {
            retHit = curHit;
            retHit.isHit = 1;
        }
    }



    if(retHit.isHit == 1){
        //计算反射方向
        retHit.reflectDir = ray.dir - 2 * dot(ray.dir, retHit.ndir) * retHit.ndir;
    }

    return retHit;
}

//光距离衰减率
float lightIntensityRate(float dis)
{
    float rate = 1- dis/MAX_DISTANCE;
    rate = clamp(rate,0,1);
    return rate*rate;
}

vec4 lightHit(vec3 pos)
{
    vec4 temp = SceneDataAt(0);
    float count = temp.r;
    float index = 1;

    vec4 pointColor = vec4(0);    

    for(int i = 0;i<count;i++)
    {
        MeshHead meshHead = MeshHeadAt(index);
        if(meshHead.len < 1) break;  //数据错误

        if(meshHead.type == MESH_POINTLIGHT)
        {
            //计算与光源之间是否有遮挡
            PointLight light = PointLightAt(index);
            Ray ray;
            ray.dir = normalize(light.pos - pos);
            ray.pos = pos + ray.dir * 0.0001;
            RayHit hit = worldHit(ray);
            bool isHide = false;
            if(hit.isHit == 1){
                if(length(hit.pos - pos) < length(light.pos - pos)){
                    isHide = true;
                }
            }

            if(!isHide){
                float rate = lightIntensityRate(length(light.pos - pos));
                pointColor += light.color * light.intensity *rate;
            }
        }
    
        index += meshHead.len;
    }

    return pointColor;
}



float random0(float x)
{
    float y = fract(sin(x)*1000);
    return y;
}

float random(float x)
{
    //随机数
    float r1 = random0(x);

    //随机正负
    float negtive = (int(r1*100)%2 == 0?-1:1);

    return r1 * negtive;
}

//反射角度采样器
vec3 makeReflectDir(RayHit hit,int index){

    if(hit.material == MATERIAL_REFLECT)
    {
        return hit.reflectDir;
    }
    else
    {
        float x = (index % Sample_Count+1)*1.0f/(Sample_Count+1);
        float y = (index / Sample_Count+1)*1.0f/(Sample_Count+1);

        float dis = length(hit.pos - pos);

        //生成球体反射向量
        vec3 retDir = vec3(
            random(hit.reflectDir.x + dis-index/10),
            random(hit.reflectDir.y+ dis+index/10),
            random(hit.reflectDir.z+ dis+index*7/10)
        );

        retDir = normalize(retDir)* 0.5 + hit.reflectDir * 0.5;
        if(dot(retDir,hit.ndir) < 0.1) retDir = -retDir;
        return retDir;
    }
}

//计算根据射线与碰撞点，表面颜色，光颜色，对光进行混合计算
vec3 BlendColor(RayHit surfaceInfo ,vec3 hitPos,vec3 lightColor){
    //衰减率
    float rate = lightIntensityRate(length(hitPos - surfaceInfo.pos));
    lightColor *=  rate;

    //根据材质衰减
    vec3 surfaceColor = surfaceInfo.color.rgb * 1.0;

    return surfaceColor * lightColor;
}

RayTraceStackItem initStackInfo(RayHit hit)
{
    RayTraceStackItem ret;
    ret.lightColor = vec3(0);
    ret.index = 0;
    ret.hit = hit;
    return ret;
}

//光追
vec3 RayTracing(Ray ray){
    //采用展开递归的方式，多次遍历世界
    Ray tempRay = ray;

    RayTraceStackItem stackInfo[MAX_DEPTH+1];
    float sampleWeight = 1.0 / Sample_Count / Sample_Count;
    //在一个深度要采样Sample_Count次
    const int SampleLength= Sample_Count * Sample_Count;
        
    //第一层碰撞检测
    RayHit hit = worldHit(ray);
    if(hit.isHit == 0) return vec3(0);
    
    //栈0为最初值
    stackInfo[0] = initStackInfo(hit);
    stackInfo[0].lightColor = hit.color * lightHit(hit.pos).rgb * (1-TRACE_COLOR_WEIGHT);

    int idepth = 0;

    //循环对递归进行展开
    int TraceDepth_MAX = clamp(traceDepth,0,MAX_DEPTH);

    while(true)
    {
       if(idepth < TraceDepth_MAX){
            //若不是最深层次，则需要采样
            int isample = stackInfo[idepth].index;
            for(;isample <SampleLength;isample++)
            {
                RayHit lastHit = stackInfo[idepth].hit;

                //设置下一次的位置
                stackInfo[idepth].index = isample+1;

                //构造当前的射线
                Ray curRay;
                curRay.pos = lastHit.pos;
                
                //漫反射
                vec3 refDir = makeReflectDir(lastHit,isample);

                //镜面反射
                //vec3 refDir = lastHit.reflectDir;

                curRay.dir = refDir;

                RayHit hit = worldHit(curRay);
                if(hit.isHit == 0) continue;

                //进入下一个深度，初始化下一层的信息
                stackInfo[idepth+1] = initStackInfo(hit);

                //初始为当前点在光源下的颜色
                float weight = idepth+1 >= TraceDepth_MAX?1:(1-TRACE_COLOR_WEIGHT);
                stackInfo[idepth+1].lightColor = hit.color.rgb * lightHit(hit.pos).rgb * weight;

                //进入下一级
                idepth++;
                break;
            }

            if(isample < SampleLength){
                //没有循环完毕直接进入下一层
                continue;
            }
        }

        if(idepth == 0){
            //最顶层，跳出光追
            break;
        }
       
        //完成当前层次的计算，将结果累加到上一层
        vec3 traceColor = BlendColor(
                    stackInfo[idepth-1].hit
                    ,stackInfo[idepth].hit.pos
                    ,stackInfo[idepth].lightColor
                )*sampleWeight * TRACE_COLOR_WEIGHT;
        stackInfo[idepth-1].lightColor += traceColor;

        //返回上一层
        idepth--;
        
    }
    
    //返回时，回复权重值
    return stackInfo[0].lightColor;
}


void main(void) {
    sceneDataSize = textureSize(sceneData, 0);
    out_color.a = 1.0;

    mat4 viewInv = inverse(view);
    mat4 projectionInv = inverse(projection);

    Ray ray;
    ray.pos = vec3(viewInv * vec4(vec3(0), 1.0));

    //从相机发射射线
    vec3 eyePos = vec3(inverse(projection) * vec4(pos.xy, 1.0, 1.0));
    vec4 worldPos = inverse(view) * vec4(eyePos,1);
    ray.dir = normalize(worldPos.xyz - ray.pos);
   
    out_color.rgb = RayTracing(ray);
}